{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4620bc8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3 install gpt-2-simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "477a8024",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:\n",
      "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "  * https://github.com/tensorflow/io (for I/O related ops)\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from gpt_2_simple.src import model as gpt2_model, encoder\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9528fd8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "params = '117m-hparams.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc10d606",
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = gpt2_model.default_hparams()\n",
    "with open(params) as f:\n",
    "    hparams.override_from_dict(json.load(f))\n",
    "\n",
    "with open('encoder.json', 'r') as f:\n",
    "    en = json.load(f)\n",
    "with open('vocab.bpe', 'r', encoding=\"utf-8\") as f:\n",
    "    bpe_data = f.read()\n",
    "    \n",
    "bpe_merges = [\n",
    "    tuple(merge_str.split()) for merge_str in bpe_data.split('\\n')[1:-1]\n",
    "]\n",
    "enc_malay = encoder.Encoder(encoder=en, bpe_merges=bpe_merges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "05b35df4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-5-27e439863b60>:25: map_and_batch (from tensorflow.contrib.data.python.ops.batching) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.experimental.map_and_batch(...)`.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow_core/contrib/data/python/ops/batching.py:276: map_and_batch (from tensorflow.python.data.experimental.ops.batching) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Dataset.map(map_func, num_parallel_calls)` followed by `tf.data.Dataset.batch(batch_size, drop_remainder)`. Static tf.data optimizations will take care of using the fused implementation.\n",
      "WARNING:tensorflow:Entity <function <lambda> at 0x7f551ac63d90> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: module 'gast' has no attribute 'Str'\n",
      "WARNING: Entity <function <lambda> at 0x7f551ac63d90> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: module 'gast' has no attribute 'Str'\n",
      "WARNING:tensorflow:From <ipython-input-5-27e439863b60>:11: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.cast` instead.\n",
      "WARNING:tensorflow:From <ipython-input-5-27e439863b60>:28: DatasetV1.make_one_shot_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_one_shot_iterator(dataset)`.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': <tf.Tensor 'IteratorGetNext:0' shape=(2, 1024) dtype=int32>}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size = 2\n",
    "maxlen = 1024\n",
    "num_cpu_threads = 1\n",
    "\n",
    "def _decode_record(record, name_to_features):\n",
    "    \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n",
    "    example = tf.parse_single_example(record, name_to_features)\n",
    "    for name in list(example.keys()):\n",
    "        t = example[name]\n",
    "        if t.dtype == tf.int64:\n",
    "            t = tf.to_int32(t)\n",
    "        example[name] = t\n",
    "\n",
    "    return example\n",
    "\n",
    "name_to_features = {\n",
    "    'input_ids': tf.FixedLenFeature([1024], tf.int64)\n",
    "}\n",
    "d = tf.data.TFRecordDataset(['gs://mesolitica-tpu-general/gpt2-testset/karangan.tfrecord'])\n",
    "d = d.apply(\n",
    "    tf.contrib.data.map_and_batch(\n",
    "        lambda record: _decode_record(record, name_to_features),\n",
    "        batch_size=batch_size,\n",
    "        num_parallel_batches=num_cpu_threads,\n",
    "        drop_remainder=True,\n",
    "    )\n",
    ")\n",
    "d = d.make_one_shot_iterator().get_next()\n",
    "d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cd7b25a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'Mean:0' shape=() dtype=float32>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = d['input_ids']\n",
    "output = gpt2_model.model(hparams=hparams, X=input_ids)\n",
    "loss = tf.reduce_mean(\n",
    "    input_tensor=tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "        labels=input_ids[:, 1:], logits=output['logits'][:, :-1]\n",
    "    )\n",
    ")\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c998ebaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f751b579",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from gs://mesolitica-tpu-general/gpt2-117m/model.ckpt-435300\n"
     ]
    }
   ],
   "source": [
    "var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)\n",
    "saver = tf.train.Saver(var_list = var_list)\n",
    "saver.restore(sess, 'gs://mesolitica-tpu-general/gpt2-117m/model.ckpt-435300')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fa944274",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.4431666\n",
      "1.5809603\n",
      "1.7055348\n",
      "1.8630769\n",
      "1.6984978\n",
      "1.767755\n",
      "1.8292981\n",
      "1.7617803\n",
      "1.8705689\n",
      "1.731211\n",
      "1.8388227\n",
      "2.0299616\n",
      "2.1151528\n",
      "1.8211054\n",
      "1.9403484\n",
      "1.9824086\n",
      "1.8927426\n",
      "1.8779304\n",
      "1.9722944\n",
      "1.8426896\n",
      "1.6888677\n",
      "1.6522177\n",
      "1.6237339\n",
      "1.771958\n",
      "1.9789941\n",
      "2.311753\n",
      "1.9610463\n",
      "1.9511932\n",
      "1.9196131\n",
      "1.994197\n",
      "1.8725495\n",
      "1.6075294\n",
      "1.8378532\n",
      "2.0257928\n",
      "1.9228168\n",
      "1.7626024\n",
      "1.7001889\n",
      "1.6126157\n",
      "1.678913\n",
      "1.8125083\n",
      "1.7683703\n",
      "End of sequence\n",
      "\t [[node IteratorGetNext (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py:1748) ]]\n",
      "\n",
      "Original stack trace for 'IteratorGetNext':\n",
      "  File \"/usr/lib/python3.6/runpy.py\", line 193, in _run_module_as_main\n",
      "    \"__main__\", mod_spec)\n",
      "  File \"/usr/lib/python3.6/runpy.py\", line 85, in _run_code\n",
      "    exec(code, run_globals)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py\", line 16, in <module>\n",
      "    app.launch_new_instance()\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/traitlets/config/application.py\", line 664, in launch_instance\n",
      "    app.start()\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/kernelapp.py\", line 619, in start\n",
      "    self.io_loop.start()\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/platform/asyncio.py\", line 199, in start\n",
      "    self.asyncio_loop.run_forever()\n",
      "  File \"/usr/lib/python3.6/asyncio/base_events.py\", line 438, in run_forever\n",
      "    self._run_once()\n",
      "  File \"/usr/lib/python3.6/asyncio/base_events.py\", line 1451, in _run_once\n",
      "    handle._run()\n",
      "  File \"/usr/lib/python3.6/asyncio/events.py\", line 145, in _run\n",
      "    self._callback(*self._args)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/ioloop.py\", line 688, in <lambda>\n",
      "    lambda f: self._run_callback(functools.partial(callback, future))\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/ioloop.py\", line 741, in _run_callback\n",
      "    ret = callback()\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 814, in inner\n",
      "    self.ctx_run(self.run)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 162, in _fake_ctx_run\n",
      "    return f(*args, **kw)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 775, in run\n",
      "    yielded = self.gen.send(value)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\", line 374, in dispatch_queue\n",
      "    yield self.process_one()\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 250, in wrapper\n",
      "    runner = Runner(ctx_run, result, future, yielded)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 741, in __init__\n",
      "    self.ctx_run(self.run)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 162, in _fake_ctx_run\n",
      "    return f(*args, **kw)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 775, in run\n",
      "    yielded = self.gen.send(value)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\", line 358, in process_one\n",
      "    yield gen.maybe_future(dispatch(*args))\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 234, in wrapper\n",
      "    yielded = ctx_run(next, result)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 162, in _fake_ctx_run\n",
      "    return f(*args, **kw)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\", line 261, in dispatch_shell\n",
      "    yield gen.maybe_future(handler(stream, idents, msg))\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 234, in wrapper\n",
      "    yielded = ctx_run(next, result)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 162, in _fake_ctx_run\n",
      "    return f(*args, **kw)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\", line 538, in execute_request\n",
      "    user_expressions, allow_stdin,\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 234, in wrapper\n",
      "    yielded = ctx_run(next, result)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tornado/gen.py\", line 162, in _fake_ctx_run\n",
      "    return f(*args, **kw)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/ipkernel.py\", line 302, in do_execute\n",
      "    res = shell.run_cell(code, store_history=store_history, silent=silent)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/ipykernel/zmqshell.py\", line 539, in run_cell\n",
      "    return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/IPython/core/interactiveshell.py\", line 2867, in run_cell\n",
      "    raw_cell, store_history, silent, shell_futures)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/IPython/core/interactiveshell.py\", line 2895, in _run_cell\n",
      "    return runner(coro)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/IPython/core/async_helpers.py\", line 68, in _pseudo_sync_runner\n",
      "    coro.send(None)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/IPython/core/interactiveshell.py\", line 3072, in run_cell_async\n",
      "    interactivity=interactivity, compiler=compiler, result=result)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/IPython/core/interactiveshell.py\", line 3263, in run_ast_nodes\n",
      "    if (await self.run_code(code, result,  async_=asy)):\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/IPython/core/interactiveshell.py\", line 3343, in run_code\n",
      "    exec(code_obj, self.user_global_ns, self.user_ns)\n",
      "  File \"<ipython-input-5-27e439863b60>\", line 28, in <module>\n",
      "    d = d.make_one_shot_iterator().get_next()\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/data/ops/iterator_ops.py\", line 426, in get_next\n",
      "    name=name)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/gen_dataset_ops.py\", line 2518, in iterator_get_next\n",
      "    output_shapes=output_shapes, name=name)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/op_def_library.py\", line 794, in _apply_op_helper\n",
      "    op_def=op_def)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/util/deprecation.py\", line 507, in new_func\n",
      "    return func(*args, **kwargs)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py\", line 3357, in create_op\n",
      "    attrs, op_def, compute_device)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py\", line 3426, in _create_op_internal\n",
      "    op_def=op_def)\n",
      "  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py\", line 1748, in __init__\n",
      "    self._traceback = tf_stack.extract_stack()\n",
      "\n"
     ]
    }
   ],
   "source": [
    "losses = []\n",
    "while True:\n",
    "    try:\n",
    "        l = sess.run(loss)\n",
    "        print(l)\n",
    "        losses.append(l)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a9cfe89a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.232461"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.exp(np.mean(losses))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4f5b029",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
